import time
import os
from transformers import AutoModel
from PIL import Image
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from tqdm import tqdm
from typing import List, Dict

# 全局变量，用于存储已经初始化好的组件
extractor = None
analyzer = None

# 模型初始化（单例模式）
def init_model(model_name="facebook/dinov2-base"):
    model = AutoModel.from_pretrained(model_name)
    return model.eval()

# 图像预处理管道
def create_preprocess_pipeline(augment=False):
    base_transforms = [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
    
    if augment:
        base_transforms.insert(1, transforms.RandomHorizontalFlip())
        base_transforms.insert(1, transforms.ColorJitter(0.4, 0.4, 0.4, 0.1))
    
    return transforms.Compose(base_transforms)

# 特征提取器（支持测试/样本不同处理）
class FeatureExtractor:
    def __init__(self, num_augments=5):
        self.model = init_model()
        self.num_augments = num_augments
        self.sample_preprocess = create_preprocess_pipeline(augment=True)  # 样本增强
        self.test_preprocess = create_preprocess_pipeline(augment=False)   # 测试不增强

    def load_image(self, image_path: str, is_test: bool) -> List[torch.Tensor]:
        try:
            image = Image.open(image_path).convert('RGB')
            preprocessor = self.test_preprocess if is_test else self.sample_preprocess
            num = 1 if is_test else self.num_augments
            return [preprocessor(image) for _ in range(num)]
        except Exception as e:
            print(f"Error loading {image_path}: {str(e)}")
            return []

    def get_features(self, image_path: str, is_test: bool) -> torch.Tensor:
        images = self.load_image(image_path, is_test)
        if not images:
            return None
            
        with torch.no_grad():
            features = torch.cat([
                self.model(img.unsqueeze(0)).last_hidden_state[:, 0, :]
                for img in images
            ])
            
        return features.mean(dim=0, keepdim=True)

# 改进后的特征银行（支持路径识别缓存）
class FeatureBank:
    _pos_cache: Dict[str, torch.Tensor] = {}  # 路径到特征的映射
    _neg_cache: Dict[str, torch.Tensor] = {}

    @staticmethod
    def initialize_cache(positive_folder: str, negative_folder: str, extractor: FeatureExtractor):
        # 只在路径未缓存时初始化
        if positive_folder not in FeatureBank._pos_cache:
            FeatureBank._pos_cache[positive_folder] = FeatureBank.batch_extract(
                extractor, positive_folder, 'positive')
            
        if negative_folder not in FeatureBank._neg_cache:
            FeatureBank._neg_cache[negative_folder] = FeatureBank.batch_extract(
                extractor, negative_folder, 'negative')

    @staticmethod
    def batch_extract(extractor: FeatureExtractor, folder_path: str, cache_key: str) -> torch.Tensor:
        # 根据缓存类型选择存储位置
        cache = FeatureBank._pos_cache if cache_key == 'positive' else FeatureBank._neg_cache
        
        # 已缓存直接返回
        if folder_path in cache:
            return cache[folder_path]

        image_paths = [os.path.join(folder_path, f) 
                      for f in os.listdir(folder_path) 
                      if f.endswith(('.png', '.jpg', '.jpeg'))]
        
        features = []
        for path in tqdm(image_paths, desc=f"Processing {cache_key} images"):
            # 样本图像使用增强处理
            feat = extractor.get_features(path, is_test=False)
            if feat is not None:
                features.append(feat)
        
        if not features:
            return None
        
        # 计算平均特征并缓存
        avg_features = torch.cat(features).mean(dim=0, keepdim=True)
        cache[folder_path] = avg_features
        return avg_features

# 相似度分析器（保持不变）
class SimilarityAnalyzer:
    def __init__(self, temperature=0.07, min_confidence=0.65):
        self.temperature = temperature
        self.min_confidence = min_confidence

    def __call__(self, test_feat: torch.Tensor, pos_feat: torch.Tensor, neg_feat: torch.Tensor):
        sim_pos = F.cosine_similarity(test_feat, pos_feat)
        sim_neg = F.cosine_similarity(test_feat, neg_feat)
        
        similarities = torch.stack([sim_pos, sim_neg]) / self.temperature
        probs = F.softmax(similarities, dim=0)
        
        return {
            "positive": (sim_pos.item(), probs[0].item()),
            "negative": (sim_neg.item(), probs[1].item())
        }

# 结果判断（保持不变）
def make_judgement(results, min_confidence):
    pos_prob = results["positive"][1]
    neg_prob = results["negative"][1]
    
    confidence = max(pos_prob, neg_prob)
    if confidence < min_confidence:
        return "不确定结果 (置信度不足)"
    
    if pos_prob > neg_prob:
        return f"产品OK (置信度: {pos_prob:.2%})"
    return f"产品NG (置信度: {neg_prob:.2%})"

def check(obj_path: str, 
         positive_folder: str, 
         negative_folder: str,
         num_augments: int = 5,
         temperature: float = 0.07,
         min_confidence: float = 0.65):
    
    global extractor, analyzer
    
    start_time = time.time()
    
    # 初始化组件（保持不变）
    if extractor is None or analyzer is None:
        print("初始化组件...")
        extractor = FeatureExtractor(num_augments=num_augments)
        analyzer = SimilarityAnalyzer(temperature, min_confidence)
    
    # 初始化路径感知缓存
    print("初始化缓存...")
    cache_start = time.time()
    FeatureBank.initialize_cache(positive_folder, negative_folder, extractor)
    cache_end = time.time()
    print(f"初始化缓存耗时: {cache_end - cache_start:.2f}秒")
    
    # 提取测试特征（使用无增强处理）
    print("提取测试图像特征...")
    test_feature_start = time.time()
    test_feature = extractor.get_features(obj_path, is_test=True)
    test_feature_end = time.time()
    print(f"提取测试图像特征耗时: {test_feature_end - test_feature_start:.2f}秒")
    
    # 获取缓存特征
    pos_features = FeatureBank._pos_cache[positive_folder]
    neg_features = FeatureBank._neg_cache[negative_folder]
    
    # 验证特征有效性
    if test_feature is None:
        raise ValueError("测试图像特征提取失败")
    if pos_features is None:
        raise ValueError("正样本特征提取失败")
    if neg_features is None:
        raise ValueError("负样本特征提取失败")
    
    # 分析结果
    results = analyzer(test_feature, pos_features, neg_features)
    judgement = make_judgement(results, min_confidence)
    
    # 打印输出（保持不变）
    print("\n===== 检测结果 =====")
    print(f"正样本相似度: {results['positive'][0]:.4f} → OK概率: {results['positive'][1]:.2%}")
    print(f"负样本相似度: {results['negative'][0]:.4f} → NG概率: {results['negative'][1]:.2%}")
    print(f"\n最终判定: {judgement}")
    print(f"总耗时: {time.time() - start_time:.2f}秒")

if __name__ == "__main__":
    # 使用示例
    check(
        obj_path='test_image2.jpg',
        positive_folder='positive_samples',
        negative_folder='negative_samples',
        num_augments=5,
        temperature=0.05,
        min_confidence=0.6
    )
    check('test_image2.jpg','positive_samples','negative_samples',5,0.05,0.6)

    check('1.jpg','positive_samples2','negative_samples2',5,0.05,0.6)
